#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
OpenMM 分子动力学模拟脚本
支持 Amber 和 GROMACS 格式文件
"""

import os
import sys
import argparse
from openmm.app import *
from openmm import *
from openmm.unit import *
import numpy as np


# Define a subclass of Minimization Reporter
class Reporter(MinimizationReporter):
    interval = 10  # report interval
    energies = []  # array to record progress

    def report(self, iteration, x, grad, args):
        # print current system energy to screen
        if iteration % self.interval == 0:
            print(iteration, args["system energy"])

        # save energy at each iteration to an array we can use later
        self.energies.append(args["system energy"])

        # The report method must return a bool specifying if minimization should be stopped.
        # You can use this functionality for early termination.
        return False


class MDSimulation:
    def __init__(self):
        self.system = None
        self.topology = None
        self.positions = None
        self.simulation = None

    def load_amber_files(self, prmtop_file, inpcrd_file):
        """加载 Amber 格式文件"""
        print(f"加载 Amber 文件: {prmtop_file}, {inpcrd_file}")

        # 加载参数文件和坐标文件
        prmtop = AmberPrmtopFile(prmtop_file)
        inpcrd = AmberInpcrdFile(inpcrd_file)

        self.system = prmtop.createSystem(
            nonbondedMethod=PME, nonbondedCutoff=1.0 * nanometer, constraints=HBonds
        )
        self.topology = prmtop.topology
        self.positions = inpcrd.positions

        # 如果有盒子信息，设置周期性边界条件
        if inpcrd.boxVectors is not None:
            self.system.setDefaultPeriodicBoxVectors(*inpcrd.boxVectors)

    def load_gromacs_files(self, top_file, gro_file):
        """加载 GROMACS 格式文件"""
        print(f"加载 GROMACS 文件: {top_file}, {gro_file}")

        # 加载拓扑文件和坐标文件
        gro = GromacsGroFile(gro_file)
        top = GromacsTopFile(top_file, periodicBoxVectors=gro.getPeriodicBoxVectors())

        self.system = top.createSystem(
            nonbondedMethod=PME, nonbondedCutoff=1.0 * nanometer, constraints=HBonds
        )
        self.topology = top.topology
        self.positions = gro.positions

    def setup_simulation(
        self,
        platform="CUDA",
        temperature=300 * kelvin,
        friction=1 / picosecond,
        timestep=2 * femtoseconds,
    ):
        """设置模拟参数"""
        print("设置模拟参数...")

        # 添加温度耦合
        self.system.addForce(AndersenThermostat(temperature, friction))

        # 选择计算平台
        try:
            if platform == "CUDA":
                platform = Platform.getPlatformByName("CUDA")
                properties = {"CudaPrecision": "mixed"}
            elif platform == "OpenCL":
                platform = Platform.getPlatformByName("OpenCL")
                properties = {"OpenCLPrecision": "mixed"}
            elif platform == "CPU":
                platform = Platform.getPlatformByName("CPU")
                properties = {}
            else:
                platform = Platform.getPlatformByName("Reference")
                properties = {}
        except Exception as e:
            print(f"无法使用 {platform} 平台，切换到 CPU 平台")
            platform = Platform.getPlatformByName("CPU")
            properties = {}

        # 创建积分器
        integrator = LangevinMiddleIntegrator(temperature, friction, timestep)

        # 创建模拟对象
        self.simulation = Simulation(
            self.topology, self.system, integrator, platform, properties
        )
        self.simulation.context.setPositions(self.positions)

        print(f"使用平台: {platform.getName()}")

    def minimize_energy(self, max_iterations=1000):
        """能量最小化"""
        print("进行能量最小化...")

        # 获取最小化前的能量
        state = self.simulation.context.getState(getEnergy=True)
        initial_energy = state.getPotentialEnergy()
        print(f"初始势能: {initial_energy}")
        reporter = Reporter()
        # 能量最小化
        self.simulation.minimizeEnergy(reporter=reporter)

        # 获取最小化后的能量
        state = self.simulation.context.getState(getEnergy=True)
        final_energy = state.getPotentialEnergy()
        print(f"最小化后势能: {final_energy}")

    def equilibration(self, steps=10000, output_freq=1000):
        """平衡阶段"""
        print(f"开始平衡模拟 ({steps} 步)...")

        # 设置初始速度
        self.simulation.context.setVelocitiesToTemperature(300 * kelvin)

        # 运行平衡
        for i in range(0, steps, output_freq):
            self.simulation.step(output_freq)
            state = self.simulation.context.getState(getEnergy=True, getPositions=True)

            potential_energy = state.getPotentialEnergy()
            kinetic_energy = state.getKineticEnergy()
            total_energy = potential_energy + kinetic_energy
            temperature = (
                2 * kinetic_energy / (3 * len(self.positions) * BOLTZMANN_CONSTANT_kB)
            ).value_in_unit(kelvin)

            print(
                f"步数: {i+output_freq}, 温度: {temperature:.2f} K, "
                f"势能: {potential_energy:.2f}, 动能: {kinetic_energy:.2f}"
            )

    def production_run(
        self,
        steps=100000,
        output_freq=1000,
        traj_file="trajectory.xtc",
        state_file="state.csv",
    ):
        """生产阶段模拟"""
        print(f"开始生产模拟 ({steps} 步)...")

        # 设置轨迹输出
        self.simulation.reporters.append(XTCReporter(traj_file, output_freq))

        # 设置状态输出
        self.simulation.reporters.append(
            StateDataReporter(
                state_file,
                output_freq,
                step=True,
                time=True,
                potentialEnergy=True,
                kineticEnergy=True,
                totalEnergy=True,
                temperature=True,
                volume=True,
                density=True,
            )
        )

        # 设置标准输出
        self.simulation.reporters.append(
            StateDataReporter(
                sys.stdout,
                output_freq,
                step=True,
                time=True,
                potentialEnergy=True,
                temperature=True,
                totalEnergy=True,
                progress=True,
                remainingTime=True,
                totalSteps=steps,
                speed=True,
            )
        )

        # 运行生产模拟
        self.simulation.step(steps)

        # 保存最终状态
        state = self.simulation.context.getState(getPositions=True, getVelocities=True)
        with open("final_state.xml", "w") as f:
            f.write(XmlSerializer.serialize(state))

        print("模拟完成！")


def main():
    parser = argparse.ArgumentParser(description="OpenMM 分子动力学模拟脚本")
    parser.add_argument(
        "-f",
        "--format",
        choices=["amber", "gromacs"],
        required=True,
        help="输入文件格式 (amber 或 gromacs)",
    )
    parser.add_argument(
        "-p", "--top", required=True, help="拓扑文件 (Amber: .prmtop, GROMACS: .top)"
    )
    parser.add_argument(
        "-c",
        "--coord",
        required=True,
        help="坐标文件 (Amber: .inpcrd/.rst7, GROMACS: .gro)",
    )
    parser.add_argument(
        "--platform",
        default="CUDA",
        choices=["CUDA", "OpenCL", "CPU", "Reference"],
        help="计算平台 (默认: CUDA)",
    )
    parser.add_argument(
        "-t", "--temperature", type=float, default=300.0, help="模拟温度 (K, 默认: 300)"
    )
    parser.add_argument(
        "-s",
        "--steps",
        type=int,
        default=50000000,
        help="生产模拟步数 (默认: 50000000)",
    )
    parser.add_argument(
        "--eq-steps", type=int, default=1000, help="平衡模拟步数 (默认: 1000)"
    )
    parser.add_argument(
        "--output-freq", type=int, default=500000, help="输出频率 (默认: 500000)"
    )
    parser.add_argument(
        "--traj", default="trajectory.xtc", help="轨迹文件名 (默认: trajectory.xtc)"
    )
    parser.add_argument(
        "-sk-min", "--skip-minimization", action="store_true", help="跳过能量最小化"
    )
    parser.add_argument(
        "-sk-eq", "--skip-equilibration", action="store_true", help="跳过平衡阶段"
    )

    args = parser.parse_args()

    # 检查输入文件是否存在
    if not os.path.exists(args.top):
        print(f"错误: 拓扑文件 {args.top} 不存在")
        sys.exit(1)
    if not os.path.exists(args.coord):
        print(f"错误: 坐标文件 {args.coord} 不存在")
        sys.exit(1)

    # 创建模拟对象
    md = MDSimulation()

    try:
        # 加载文件
        if args.format == "amber":
            md.load_amber_files(args.top, args.coord)
        elif args.format == "gromacs":
            md.load_gromacs_files(args.top, args.coord)

        # 设置模拟
        md.setup_simulation(
            platform=args.platform, temperature=args.temperature * kelvin
        )

        # 能量最小化
        if not args.skip_minimization:
            md.minimize_energy()

        # 平衡阶段
        if not args.skip_equilibration:
            md.equilibration(steps=args.eq_steps, output_freq=args.output_freq)

        # 生产模拟
        md.production_run(
            steps=args.steps, output_freq=args.output_freq, traj_file=args.traj
        )

    except Exception as e:
        print(f"模拟过程中发生错误: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    main()
